from tornado import gen
from tornado import httpclient
from tornado import httputil
from tornado import ioloop
from tornado import websocket

import functools
import json
import time

from HDetection import create_model
from HDetectionv3 import hdetect
from LDetection2 import ldetect2
import json
import cv2
# Social Distance relate
import AI_Algorithims.SocialDisAnalyser.analyser as analyser
from yolo.tool.utils import load_class_names
from yolo.tool.torch_utils import do_detect
from Websocket.tornado_websocket_server import LandArg, SocialDisArg

APPLICATION_JSON = 'application/json'

DEFAULT_CONNECT_TIMEOUT = 30
DEFAULT_REQUEST_TIMEOUT = 30
SDArgs = SocialDisArg(angle_factor=0.8, H_zoom_factor=1.2, confid=0.3)
arg_single = LandArg(True, 544, 3, 0.1, 0.3, "")

def is_json(myjson):
    try:
        json_object = json.loads(myjson)
    except ValueError as e:
        return False
    return True


class WebSocketClient(object):
    """Base for web socket clients.
    """

    DISCONNECTED = 0
    CONNECTING = 1
    CONNECTED = 2

    def __init__(self, io_loop=None,
                 connect_timeout=DEFAULT_CONNECT_TIMEOUT,
                 request_timeout=DEFAULT_REQUEST_TIMEOUT):

        self.connect_timeout = connect_timeout
        self.request_timeout = request_timeout
        self._io_loop = io_loop or ioloop.IOLoop.current()
        self._ws_connection = None
        self._connect_status = self.DISCONNECTED

    def connect(self, ws_url):
        """Connect to the server.
        :param str ws_url: server URL.
        """
        self._connect_status = self.CONNECTING
        headers = httputil.HTTPHeaders({'Content-Type': APPLICATION_JSON})
        request = httpclient.HTTPRequest(url=ws_url,
                                         connect_timeout=self.connect_timeout,
                                         request_timeout=self.request_timeout,
                                         headers=headers)
        ws_conn = websocket.WebSocketClientConnection(request, self._io_loop)
        ws_conn.connect_future.add_done_callback(self._connect_callback)

    def send(self, data):
        """Send message to the server
        :param str data: message.
        """

        if self._ws_connection:
            self._ws_connection.write_message(str(data))

    def close(self, reason=''):
        """Close connection.
        """

        if self._connect_status != self.DISCONNECTED:
            self._connect_status = self.DISCONNECTED
            self._ws_connection and self._ws_connection.close()
            self._ws_connection = None
            self.on_connection_close(reason)

    def _connect_callback(self, future):
        if future.exception() is None:
            self._connect_status = self.CONNECTED
            self._ws_connection = future.result()
            self.on_connection_success()
            self._read_messages()
        else:
            self.close(future.exception())

    def is_connected(self):
        return self._ws_connection is not None

    @gen.coroutine
    def _read_messages(self):
        while True:
            msg = yield self._ws_connection.read_message()
            if msg is None:
                self.close()
                break

            self.on_message(msg)

    def on_message(self, msg):
        """This is called when new message is available from the server.
        :param str msg: server message.
        """

        pass

    def on_connection_success(self):
        """This is called on successful connection ot the server.
        """

        pass

    def on_connection_close(self, reason):
        """This is called when server closed the connection.
        """
        pass


class RTCWebSocketClient(WebSocketClient):
    msg = {'cmd': 'msg', 'from': 'Frankie',
           'to': 'Peter', 'body': 'Hello, Peter'}
    join_msg = {'cmd': 'join', 'roomId': 'MEC', 'uid': 'MecClient'}
    hb_msg = {'cmd': 'hb'}  # hearbeat

    heartbeat_interval_in_secs = 10

    def __init__(self, io_loop=None,
                 connect_timeout=DEFAULT_CONNECT_TIMEOUT,
                 request_timeout=DEFAULT_REQUEST_TIMEOUT):

        self.connect_timeout = connect_timeout
        self.request_timeout = request_timeout
        self._io_loop = io_loop or ioloop.IOLoop.current()
        self.ws_url = None
        self.auto_reconnet = False

        self.model = create_model()
        self.classnames = load_class_names("D:/PYProject/FlaskServer/yolo/data/mark.names")

        self.curr_height = 1000
        self.rtmp_thread = ""
        self.rtmp_server = 'rtmp://8.134.15.236:1935/live/'
        self.mission_type = "land"

        self.net = analyser.getSocialDisNet()
        self.labels = analyser.getLabels()
        self.device = analyser.getDevice()

        super(RTCWebSocketClient, self).__init__(self._io_loop,
                                                 self.connect_timeout,
                                                 self.request_timeout)

    def connect(self, url, auto_reconnet=True, reconnet_interval=10):
        self.ws_url = url
        self.auto_reconnet = auto_reconnet
        self.reconnect_interval = reconnet_interval

        super(RTCWebSocketClient, self).connect(self.ws_url)

    def on_message(self, message):
        print('on_message msg=', message)
        print(len(str(message)))
        value = message
        print(value)
        if is_json(value):
            j_value = json.loads(value)
            try:
                if j_value["cmd"] == "mecdata":
                    self.cap = cv2.VideoCapture(self.rtmp_server + j_value["recordName"])
                    # 调用cv2方法获取cap的视频帧（帧：每秒多少张图片）
                    self.fps = self.cap.get(cv2.CAP_PROP_FPS)
                    print(self.fps)
                    # 获取cap视频流的每帧大小
                    self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                    self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                    self.half_width = self.width / 2
                    self.half_height = self.height / 2
                    self.size = (self.width, self.height)
                    print(self.size)

                    ret, image = self.cap.read()
                    count = 0
                    (H, W) = image.shape[:2]

                    while ret:
                        if not ret:
                            break
                        if count % 6 == 0 or count == 0:
                            print(count)
                            ###Social DIS Block###------START
                            image, rect_list, line_list = analyser.socialDisDetecter(self.net, self.labels, image,
                                                                                     self.device, H, W, SDArgs)
                            rect_list = [i for item in rect_list for i in item]
                            line_list = [i for item in line_list for i in item]
                            if len(rect_list) == 0:
                                rect_list = [0]
                            if len(line_list) == 0:
                                line_list = [0]
                            msg = {"cmd": "mecdata", "remoteRoomId": j_value["roomId"], "remoteUid": j_value["uid"], "uid": "MecClient",
                                   "type": 1, "msg": "SD", "time": 1, "deviation": count,
                                   "rect_list": str(rect_list), "line_list": str(line_list)}
                            print(msg)
                            self.send(msg)
                            ###Social DIS Block###------END

                            ###CV Landing Block###------START
                            sized = cv2.resize(image, (arg_single.size, arg_single.size))
                            sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
                            start = time.clock()
                            boxes = do_detect(self.model, sized, arg_single.conf, arg_single.nms,
                                              arg_single.use_cuda)
                            end = time.clock()
                            print((start - end), "---------------------------------------------")
                            if len(boxes[0]) == 4 or len(boxes[0]) == 5 or len(boxes[0]) == 6:
                                cmd, image, dx, dy, angle = ldetect2(boxes, image, self.half_width,
                                                                     self.half_height)
                            else:
                                angle = 0
                                cmd, image, dx, dy = hdetect(boxes, image=image, mclass_names=self.classnames)
                            if cmd == "null":
                                msg = {"cmd": "mecdata", "remoteRoomId": j_value["roomId"], "remoteUid": j_value["uid"], "uid": "MecClient",
                                       "type": 2, "msg": cmd, "time": 1, "deviation": count,
                                       "rect_list": str(rect_list), "line_list": str(line_list)}
                                print(msg)
                                self.send(msg)
                            else:
                                value = str({"dx": dx, "dy": dy, "mpitch": round(dy / self.half_height, 2),
                                             "mroll": round(dx / self.half_width, 2), "myaw": angle,
                                             "mthrottle": 0})
                                print(value)
                                msg = {"cmd": "mecdata", "remoteRoomId": j_value["roomId"], "remoteUid": j_value["uid"], "uid": "MecClient",
                                       "type": 2, "msg": value, "time": 1, "deviation": count,
                                       "rect_list": str(rect_list), "line_list": str(line_list)}
                                print(msg)
                                self.send(msg)
                            print("before", dx, dy)
                            ###CV Landing Block###------END

                            cv2.namedWindow("Landing", 0)
                            cv2.imshow("Landing", image)
                            cv2.waitKey(1)
                        count = count + 1
                        ret, image = self.cap.read()

            except ValueError as e:
                print(e)
        self._io_loop.call_later(self.heartbeat_interval_in_secs,
                                 functools.partial(self.send, self.hb_msg))

    def on_connection_success(self):
        print('Connected!')
        self.send(self.join_msg)

    def on_connection_close(self, reason):
        print('Connection closed reason=%s' % (reason,))
        self.reconnect()

    def reconnect(self):
        print('reconnect')
        if not self.is_connected() and self.auto_reconnet:
            self._io_loop.call_later(self.reconnect_interval,
                                     super(RTCWebSocketClient, self).connect, self.ws_url)


def main():
    io_loop = ioloop.IOLoop.instance()

    client = RTCWebSocketClient(io_loop)
    # ws_url = 'ws://127.0.0.1:8090/ws'
    ws_url = 'ws://8.134.15.236:3000'
    client.connect(ws_url, auto_reconnet=True, reconnet_interval=10)

    try:
        io_loop.start()
    except KeyboardInterrupt:
        client.close()


if __name__ == '__main__':
    main()
